import cv2
import BaseData
import onnxruntime as rt
import numpy as np


screen_rotation = False
cap = cv2.VideoCapture(0)  # 设置摄像头编号，如果只插了一个USB摄像头，基本上都是0
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)  # 设置摄像头图像宽度
cap.set(cv2.CAP_PROP_FRAME_HEIGHT,240)  # 设置摄像头图像高度
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 设置OpenCV内部的图像缓存，可以极大提高图像的实时性。
cv2.namedWindow('camera', cv2.WND_PROP_FULLSCREEN)  # 窗口全屏
cv2.setWindowProperty('camera', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)  # 窗口全屏

tag = ['0faya', 'faya']
def onnx_cls(img):
    sess = rt.InferenceSession('out_file/tudou.onnx', None)
    input_name = sess.get_inputs()[0].name
    out_name = sess.get_outputs()[0].name
    dt = BaseData.ImageData(img, size=(224, 224))
    input_data = dt.to_tensor()
    pred_onx = sess.run([out_name], {input_name: input_data})
    result = np.argmax(pred_onx[0], axis=1)[0]
	confidence = np.max(pred_onx[0], axis=1)[0]
    print('置信度:', confidence)

    return result
while not cap.isOpened():
    continue
while True:
    success, image = cap.read()

    #if not success:
    #    print("Ignoring empty camera frame.")
    #    break
    #if screen_rotation:  # 是否要旋转屏幕
    #    image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)  # 旋转屏幕
    #    image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)  # 旋转屏幕

    if success:
        success, image = cap.read()
        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
        idx = onnx_cls(image)
        print('result:' + tag[idx])
        #cv2.putText(image, tag[idx], (0, 40), cv2.FONT_HERSHEY_TRIPLEX, 1, (150, 0, 180), 1)
        #cv2.imshow('result',image)
		if  tag[idx] == '0faya':
			print("土豆没有发芽，可以食用！")

		else:
			print("发芽土豆有毒，不可食用！")

        cv2.imshow('camera',image)
    if cv2.waitKey(5) & 0xFF == 27:
        break
cap.release()
cv2.destroyAllWindows()
